import time
import logging
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait, Select
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException

# 设置日志
logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class VoiceEmotionMapper:
    def __init__(self):
        self.url = "https://tts.ai-lab.top"
        options = webdriver.ChromeOptions()
        options.add_argument('--headless')  # 无头模式
        options.add_argument('--no-sandbox')
        options.add_argument('--disable-dev-shm-usage')
        self.driver = webdriver.Chrome(options=options)
        
    def __enter__(self):
        return self
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.driver:
            self.driver.quit()
            
    def get_voice_emotions(self):
        """获取所有语音及其对应的感情选项"""
        try:
            # 打开网页
            logger.info("打开TTS页面")
            self.driver.get(self.url)
            time.sleep(2)  # 等待页面加载
            
            # 获取语音选择框
            voice_select = WebDriverWait(self.driver, 10).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, "select#voice"))
            )
            
            # 获取所有语音选项
            voice_options = [(option.get_attribute('value'), option.text) 
                           for option in voice_select.find_elements(By.TAG_NAME, "option")
                           if option.get_attribute('value') != '0']  # 排除"请选择语音"选项
            
            logger.info(f"找到 {len(voice_options)} 个语音选项")
            
            # 存储结果的字典
            voice_emotion_map = {}
            
            # 遍历每个语音选项
            for voice_value, voice_text in voice_options:
                logger.info(f"\n正在获取语音 {voice_text} 的感情选项")
                
                # 选择语音
                Select(voice_select).select_by_value(voice_value)
                time.sleep(3)  # 等待感情选项加载
                
                # 获取感情选择框
                emotion_select = WebDriverWait(self.driver, 10).until(
                    EC.presence_of_element_located((By.CSS_SELECTOR, "select#style"))
                )
                
                # 等待感情选项加载完成
                max_retries = 10
                retry_count = 0
                last_options_count = 0
                
                while retry_count < max_retries:
                    emotion_options = [(option.get_attribute('value'), option.text) 
                                     for option in emotion_select.find_elements(By.TAG_NAME, "option")
                                     if option.get_attribute('value') != '0']  # 排除"请选择感情"选项
                    current_count = len(emotion_options)
                    
                    if current_count > 0 and current_count == last_options_count:
                        break
                        
                    last_options_count = current_count
                    time.sleep(1)
                    retry_count += 1
                
                # 存储结果
                voice_emotion_map[voice_text] = emotion_options
                logger.info(f"语音 {voice_text} 有 {len(emotion_options)} 个感情选项")
                for emotion_value, emotion_text in emotion_options:
                    logger.info(f"  - {emotion_value}: {emotion_text}")
            
            return voice_emotion_map
        except Exception as e:
            logger.error(f"获取语音感情选项时出错: {str(e)}")
            return {}
            
    def print_voice_emotions(self):
        """打印所有语音及其感情选项"""
        voice_emotion_map = self.get_voice_emotions()
        
        print("\n=== 语音与感情选项对应关系 ===")
        for voice, emotions in voice_emotion_map.items():
            print(f"\n语音: {voice}")
            print("可用感情选项:")
            for emotion_value, emotion_text in emotions:
                print(f"  - {emotion_value}: {emotion_text}")
            print("-" * 50)
            
def main():
    with VoiceEmotionMapper() as mapper:
        mapper.print_voice_emotions()

if __name__ == "__main__":
    main() 